/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
* This file is a part of the CANN Open Software.
* Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* ===================================================================================================================*/
#include "ascir_register.h"
#include "graph/types.h"
namespace ge {
namespace ascir {
EXPORT_GENERATOR()

REG_ASC_IR_START_NODE_WITH_ATTR(Data);
REG_ASC_IR_START_NODE(Constant);
REG_ASC_IR_START_NODE(Workspace);
REG_ASC_IR_START_NODE(TbufData);
REG_ASC_IR_1IO(Output);

REG_ASC_IR_1IO(Load).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_1IO(Broadcast).UseFirstInputDataType();
REG_ASC_IR_1IO(Store).UseFirstInputDataType().UseFirstInputView();

//这里先打桩用来测试
REG_ASC_IR_1IO(WorkspaceWithInput).UseFirstInputDataType().UseFirstInputView();

/*
 * todo nop比较特别，不确定是不是缺陷，原定义中，GEIR与ASCIR是不同的，GEIR多了个必选属性
namespace ge {
REG_OP(Nop)
    .REQUIRED_ATTR(dst_type, Int)
    .INPUT(x, TensorType::ALL())
    .OUTPUT(y, TensorType::ALL())
.OP_END_FACTORY_REG(Nop)
}

namespace ascir::ops {
REG_OPS(Nop)
OPS_INPUT(0, x)
OPS_OUTPUT(0, y)
END_OPS(Nop)
}
 */
REG_ASC_IR_1IO(Nop).UseFirstInputDataType();
REG_ASC_IR_1IO(Cast);

REG_ASC_IR_1IO(Abs).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_1IO(Exp).UseFirstInputDataType().UseFirstInputView();

REG_ASC_IR_1IO(Max).UseFirstInputDataType();
REG_ASC_IR_1IO(Sum).UseFirstInputDataType();

REG_ASC_IR_2I1O(Add).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_2I1O(Sub).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_2I1O(Div).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_2I1O(Mul).UseFirstInputDataType().UseFirstInputView();

REG_ASC_IR_2I1O(GT).UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR_2I1O(Muls).UseFirstInputDataType().UseFirstInputView();

// REG_ASC_IR_2I1O(MatMul)
REG_ASC_IR(MatMul)
.Inputs({"x1", "x2"})
    .OptionalInput("x3")
    .Outputs({"y"})
    .UseFirstInputDataType();

REG_ASC_IR(FlashSoftmax)
.Inputs({"x1", "x2", "x3"})
    .Outputs({"y1", "y2", "y3"})
    .UseFirstInputDataType();
REG_ASC_IR_2I1O(Dropout).UseFirstInputDataType();
REG_ASC_IR_2I1O(Select).UseFirstInputDataType();
// 适配add_layer_norm新增的api
REG_ASC_IR(CalcMean).Inputs({"x1", "x2", "x3"}).Outputs({"y1", "y2", "y3"}).UseFirstInputDataType();
REG_ASC_IR(CalcMeanSlice).Inputs({"x1", "x2", "x3"}).Outputs({"y1", "y2", "y3"}).UseFirstInputDataType();
REG_ASC_IR(CalcRstd).Inputs({"x1", "x2", "x3"}).Outputs({"y1", "y2"}).UseFirstInputDataType();
REG_ASC_IR(CalcRstdSlice).Inputs({"x1", "x2"}).Outputs({"y1", "y2"}).UseFirstInputDataType();
REG_ASC_IR(VFWelfordPart1Update)
.Inputs({"x1", "x2", "x3"})
    .Outputs({"y1", "y2", "y3", "y4"})
    .UseFirstInputDataType().UseFirstInputView();
REG_ASC_IR(VFWelfordPart1Finalize).Inputs({"x1", "x2"}).Outputs({"y1", "y2"}).UseFirstInputDataType();
REG_ASC_IR(VFCalcYWelford).Inputs({"x1", "x2", "x3"}).Outputs({"y1"}).UseSecondInputDataType().UseFirstInputView();
REG_ASC_IR(CalcY).Inputs({"x1", "x2", "x3", "x4"}).Outputs({"y1"}).UseSecondInputDataType().UseFirstInputView();
REG_ASC_IR(CalcMeanStub)
    .Inputs({"x1", "x2", "x3"})
    .Outputs({"y1", "y2", "y3", "y4"}).Attr<int64_t>("reduce_axis_dim")
    .DataTypes({PromptDtype(0U), 0U, PromptDtype(0U),
                ge::DT_DOUBLE})
    .Views({ReduceView(0U, "reduce_axis_dim"), 0U, 0U, 0U});
}  // namespace ascir
}
